# Parquet文件系统模块 - 用于冷数据的长期存储

import pandas as pd
import os
import logging
from datetime import datetime

logger = logging.getLogger(__name__)

class ParquetStorage:
    """
    Parquet文件系统类，用于冷数据的长期存储
    实现按年/月/日的目录结构存储策略
    """
    
    def __init__(self, base_dir: str = 'data/archive'):
        """
        初始化Parquet存储系统
        
        Args:
            base_dir: 基础存储目录
        """
        try:
            self.base_dir = base_dir
            
            # 确保基础目录存在
            os.makedirs(base_dir, exist_ok=True)
            
            # 创建数据类型目录
            self.quotes_dir = os.path.join(base_dir, 'quotes')
            self.daily_dir = os.path.join(base_dir, 'daily')
            
            os.makedirs(self.quotes_dir, exist_ok=True)
            os.makedirs(self.daily_dir, exist_ok=True)
            
            logger.info(f"Parquet存储系统初始化成功: {base_dir}")
        except Exception as e:
            logger.error(f"Parquet存储系统初始化失败: {str(e)}")
            raise

    def archive_daily_data(self, df: pd.DataFrame, date: datetime = None):
        """
        归档日线数据
        
        Args:
            df: 包含日线数据的DataFrame
            date: 归档日期，默认为当前日期
        """
        try:
            if df.empty:
                logger.warning("没有日线数据需要归档")
                return
            
            if date is None:
                date = datetime.now()
            
            # 获取存储路径
            storage_dir = self._get_path_by_date(self.daily_dir, date)
            file_path = os.path.join(storage_dir, f"{date.strftime('%Y%m%d')}.parquet")
            
            # 如果文件已存在，则追加数据
            if os.path.exists(file_path):
                existing_df = pd.read_parquet(file_path)
                combined_df = pd.concat([existing_df, df]).drop_duplicates(subset=['symbol', 'date'])
                combined_df.to_parquet(file_path, index=False)
            else:
                df.to_parquet(file_path, index=False)
            
            logger.info(f"成功归档{len(df)}条日线数据到{file_path}")
        except Exception as e:
            logger.error(f"归档日线数据失败: {str(e)}")

    def _get_path_by_date(self, base_dir: str, date: datetime) -> str:
        """
        根据日期获取存储路径
        
        Args:
            base_dir: 基础目录
            date: 日期
        
        Returns:
            str: 存储路径
        """
        year_dir = os.path.join(base_dir, str(date.year))
        month_dir = os.path.join(year_dir, f"{date.month:02d}")
        os.makedirs(month_dir, exist_ok=True)
        return month_dir